package com.zhao.dota.security.store;

import com.zhao.dota.bo.AuthenticateUser;
import com.zhao.dota.security.AuthenticationKeyGenerator;
import com.zhao.dota.security.token.AccessToken;

import java.util.concurrent.ConcurrentHashMap;

public class InMemoryTokenStore implements TokenStore {

    private ConcurrentHashMap<String, AccessToken> accessTokenMap = new ConcurrentHashMap<>();

    private ConcurrentHashMap<String, String> accessTokenToRefreshTokenMap = new ConcurrentHashMap<>();

    private ConcurrentHashMap<String, String> refreshTokenToAccessTokenMap = new ConcurrentHashMap<>();

    private ConcurrentHashMap<String, AuthenticateUser> authenticateUserMap = new ConcurrentHashMap<>();

    private final ConcurrentHashMap<String, AuthenticateUser> refreshTokenAuthenticationMap = new ConcurrentHashMap<>();

    private final ConcurrentHashMap<String, AccessToken> authenticationToAccessTokenMap = new ConcurrentHashMap<>();

    private AuthenticationKeyGenerator keyGenerator = new AuthenticationKeyGenerator();


    @Override
    public AccessToken readAccessToken(String token) {
        return accessTokenMap.get(token);
    }

    @Override
    public AuthenticateUser readAuthentication(AccessToken accessToken) {
        return authenticateUserMap.get(accessToken.getValue());
    }

    @Override
    public AuthenticateUser readAuthentication(String token) {
        return authenticateUserMap.get(token);
    }

    @Override
    public void storeAccessToken(AccessToken token, AuthenticateUser authentication) {
        accessTokenMap.put(token.getValue(), token);
        authenticateUserMap.put(token.getValue(), authentication);
        accessTokenToRefreshTokenMap.put(token.getValue(), token.getRefreshToken());
        refreshTokenToAccessTokenMap.put(token.getRefreshToken(), token.getValue());
        refreshTokenAuthenticationMap.put(token.getRefreshToken(), authentication);
        String key = keyGenerator.extractKey(authentication);
        authenticationToAccessTokenMap.put(key, token);
    }

    @Override
    public void removeAccessToken(AccessToken token) {
        removeAccessToken(token.getValue());
    }

    @Override
    public void removeAccessToken(String token) {
        AccessToken accessToken = accessTokenMap.remove(token);
        AuthenticateUser authenticateUser = authenticateUserMap.remove(token);
        accessTokenToRefreshTokenMap.remove(token);
        refreshTokenToAccessTokenMap.remove(accessToken.getRefreshToken());
        refreshTokenAuthenticationMap.remove(accessToken.getRefreshToken());
        String key = keyGenerator.extractKey(authenticateUser);
        authenticationToAccessTokenMap.remove(key);
    }

    @Override
    public void storeRefreshToken(String refreshToken, AuthenticateUser authentication) {
        refreshTokenAuthenticationMap.put(refreshToken, authentication);
    }

    @Override
    public String readRefreshToken(String tokenValue) {
        return tokenValue;
    }

    @Override
    public AuthenticateUser readAuthenticationForRefreshToken(String refreshToken) {
        return refreshTokenAuthenticationMap.getOrDefault(refreshToken, null);
    }

    @Override
    public void removeRefreshToken(String refreshToken) {
        refreshTokenToAccessTokenMap.remove(refreshToken);
        refreshTokenAuthenticationMap.remove(refreshToken);
    }

    @Override
    public void removeAccessTokenUsingRefreshToken(String refreshToken) {
        String accessToken = refreshTokenToAccessTokenMap.remove(refreshToken);
        if (accessToken != null) {
            removeAccessToken(accessToken);
        }
    }

    @Override
    public AccessToken getAccessToken(AuthenticateUser authentication) {
        String key = keyGenerator.extractKey(authentication);
        return authenticationToAccessTokenMap.getOrDefault(key, null);
    }

    @Override
    public AccessToken getAccessToken(String refreshToken) {
        String accessToken = refreshTokenToAccessTokenMap.get(refreshToken);
        if (accessToken != null) {
            return accessTokenMap.getOrDefault(accessToken, null);
        }
        return null;
    }
}
